import json
import torch
import logging
from datetime import datetime
from transformers import AutoTokenizer, AutoModelForCausalLM

TARGET_GPU_ID = 1
SYSTEM_PROMPT = """
请从以下文章中，​仅提取作者以第一人称“我”或“本人”直接描述的亲身经历，要求用一段话概括，并完整覆盖以下所有相关维度的详细细节​（若涉及则提取，未涉及不补充），同时严格控制总字数在500字以内：

需重点提取的细节维度（必须包含）：
​日常生活细节​：具体时间、地点、行为动作、感官体验；
​家庭/朋友关联事件​：互动对象、对话内容、情绪反馈；
​职业/工作相关经历​：具体任务、关键步骤、突发状况；

关键约束：
输出必须基于原文，仅保留“我”直接参与或感知的内容，排除宏观事件、他人观察、抽象观点；
语言需简洁连贯，避免冗余，确保所有关键细节（时间、地点、对话、动作等）完整呈现；
总字数严格控制在500字以内（若原文细节过多，优先保留最核心的场景和情感）。
原文内容：{text}
请生成符合要求的总结：
"""

# 配置日志记录
def setup_logger():
    """配置日志记录器（文件+控制台）"""
    # 创建日志目录（如果不存在）
    log_dir = "train_cf/src/20250624/logs"
    import os
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)
    
    # 日志文件名（带时间戳）
    log_file = f"summarize_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
    log_path = os.path.join(log_dir, log_file)

    # 配置日志格式
    formatter = logging.Formatter(
        "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S"
    )

    # 控制台处理器
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    console_handler.setFormatter(formatter)

    # 文件处理器（记录所有级别）
    file_handler = logging.FileHandler(log_path, encoding="utf-8")
    file_handler.setLevel(logging.DEBUG)
    file_handler.setFormatter(formatter)

    # 获取根日志器并配置
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.DEBUG)
    logger.addHandler(console_handler)
    logger.addHandler(file_handler)

    logger.info("日志系统初始化完成，日志文件：%s", log_path)
    return logger

# 初始化全局日志器
logger = setup_logger()

def load_json_file(file_path):
    """加载JSON文件（带日志记录）"""
    try:
        logger.debug(f"尝试加载JSON文件：{file_path}")
        with open(file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        logger.info(f"成功加载JSON文件，文件大小：{len(str(data))}字符")
        return data
    except FileNotFoundError:
        logger.error(f"文件未找到：{file_path}", exc_info=True)
        raise
    except json.JSONDecodeError as e:
        logger.error(f"JSON解析失败：{str(e)}（位置：{e.pos}行：{e.lineno}列：{e.colno}）", exc_info=True)
        raise
    except Exception as e:
        logger.error(f"加载JSON文件时发生未知错误：{str(e)}", exc_info=True)
        raise

def summarize_text(text, model, tokenizer, max_length=500, device=None):
    """使用模型总结文本内容（带日志记录）"""
    try:
        full_prompt = f"{SYSTEM_PROMPT.format(text=text)}"
        # logger.debug(f"完整Prompt：{full_prompt}")  # 关键日志
        # 记录原始内容（截断至前200字符）
        original_sample = text[:200] + ("..." if len(text) > 200 else "")
        logger.info(f"开始生成摘要，原始内容示例（前200字符）：{original_sample}")
        logger.info(f"原始内容完整长度：{len(text)}字符")
        # logger.debug(f"开始生成摘要，原始文本长度：{len(text)}字符")
        inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True, max_length=500).to(device)
        
        # 生成摘要
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=max_length,
                num_beams=5,
                early_stopping=True,
                length_penalty=2,
                do_sample=True,
                temperature=0.7,
                top_p=0.7
            )
        
        summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
        summary_length = len(summary)
        logger.info(f"摘要生成完成，摘要示例：{summary}")
        logger.info(f"原始长度：{len(text)}字符 → 摘要长度：{summary_length}字符（限制{max_length}字符）")
        return summary
    except Exception as e:
        logger.error(f"文本摘要生成失败：{str(e)}", exc_info=True)
        raise

def process_json_with_model(json_data, model, tokenizer, device=None):
    """处理JSON数据（带日志记录）"""
    try:
        logger.info("开始处理JSON数据...")
        if isinstance(json_data, list):
            total = len(json_data)
            logger.debug(f"检测到列表类型数据，包含{total}个元素")
            for idx, item in enumerate(json_data, 1):
                current_msg = f"处理列表第{idx}/{total}条'）"
                logger.info(current_msg)
                if "output" in item and isinstance(item["output"], str):
                    logger.debug(f"处理第{idx+1}个元素的'output'字段")
                    item["summary"] = summarize_text(item["output"], model, tokenizer, device=device)
        elif isinstance(json_data, dict):
            logger.debug(f"检测到字典类型数据，包含{len(json_data)}个键")
            for key, value in json_data.items():
                if isinstance(value, (dict, list)):
                    logger.debug(f"递归处理键'{key}'的嵌套结构")
                    process_json_with_model(value, model, tokenizer)
                elif key == "output" and isinstance(value, str):
                    logger.debug(f"处理顶层'output'字段")
                    json_data["summary"] = summarize_text(value, model, tokenizer, device=device)
        else:
            logger.warning("不支持的JSON数据类型（非列表/字典）")
        logger.info("JSON数据处理完成")
        return json_data
    except Exception as e:
        logger.error(f"处理JSON数据时发生错误：{str(e)}", exc_info=True)
        raise

def main():
    """主函数（带日志记录）"""
    try:
        logger.info("===== 开始执行主程序 =====")

        # 1. 检测GPU可用性并设置设备
        device = torch.device(f"cuda:{TARGET_GPU_ID}" if torch.cuda.is_available() else "cpu")
        logger.info(f"检测到GPU可用：{torch.cuda.is_available()}，当前使用设备：{device}")
        if device.type == "cuda":
            logger.info(f"GPU设备信息：{torch.cuda.get_device_name(0)}")
            logger.info(f"可用GPU内存：{torch.cuda.get_device_properties(0).total_memory // (1024**3)}GB")
        
        # 2. 加载JSON文件
        # json_path = input("请输入JSON文件路径: ")
        json_path = "train_cf/data/大西瓜籽妈妈.deepseek.chat.cf_new_format.json"
        logger.info(f"JSON路径：{json_path}")
        json_data = load_json_file(json_path)
        
        # 2. 加载本地模型
        # model_path = input("请输入deepseek32b模型路径: ")
        model_path = "/data2/modelscope/hub/models/deepseek-ai/DeepSeek-R1-Distill-Qwen-32B"
        logger.info(f"模型路径：{model_path}")
        # print("正在加载模型，请稍候...")
        logger.info("开始加载大语言模型...")
        tokenizer = AutoTokenizer.from_pretrained(model_path)
        model = AutoModelForCausalLM.from_pretrained(model_path)
        model.to(device)
        logger.info("模型加载完成")
        
        # 3. 处理JSON数据
        print("正在处理数据...")
        logger.info("开始处理JSON数据中的'output'字段...")
        result = process_json_with_model(json_data, model, tokenizer, device=device)
        logger.info("数据处理完成")
        
        # 4. 保存结果
        # output_path = input("请输入保存结果的路径: ")
        output_path = "train_cf/data/summarize_output.json"
        logger.info(f"输出路径：{output_path}")
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(result, f, ensure_ascii=False, indent=2)
        logger.info(f"结果成功保存至：{output_path}")
        
        logger.info("===== 程序执行完成 =====")
    except Exception as e:
        logger.error(f"程序执行失败：{str(e)}", exc_info=True)
        print(f"程序运行出错，请查看日志文件获取详细信息")

if __name__ == "__main__":
    main()